import warnings

from mindspore import nn, ops,Tensor
from functools import partial
from mindspore import dtype as mstype
from mindspore import context
from mindspore.common.initializer import One, Normal


class UpsampleConvLayer(nn.Cell):
    def __init__(self, in_channels, out_channels, kernel_size, stride):
        super(UpsampleConvLayer, self).__init__()
        self.conv2d = nn.Conv2dTranspose(in_channels, out_channels, kernel_size, stride=stride,
                                         pad_mode='pad', padding=1, has_bias=True)

    def construct(self, x):
        out = self.conv2d(x)
        return out



class ConvLayer(nn.Cell):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding=0):
        super(ConvLayer, self).__init__()

        if padding > 0:
            self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride,
                                pad_mode='pad', padding=padding, has_bias=True)
        else:
            self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, has_bias=False)

    def construct(self, x):
        out = self.conv2d(x)
        return out


class ResidualBlock(nn.Cell):
    def __init__(self, channels):
        super(ResidualBlock, self).__init__()
        self.conv1 = ConvLayer(channels, channels, kernel_size=3, stride=1, padding=1)
        self.conv2 = ConvLayer(channels, channels, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()
        self.add = ops.Add()

    def construct(self, x):
        residual = x
        out = self.relu(self.conv1(x))
        out = self.conv2(out) * 0.1
        out = self.add(out, residual)
        return out


def resize(input,
           size=None,
           scale_factor=None,
           mode='nearest',
           align_corners=False,
           warning=True):
    if warning:
        if size is not None and align_corners:
            input_h, input_w = tuple(int(x) for x in input.shape[2:])
            output_h, output_w = tuple(int(x) for x in size)
            if output_h > input_h or output_w > output_h:
                if ((output_h > 1 and output_w > 1 and input_h > 1
                     and input_w > 1) and (output_h - 1) % (input_h - 1)
                        and (output_w - 1) % (input_w - 1)):
                    warnings.warn(
                        f'When align_corners={align_corners}, '
                        'the output would more aligned if '
                        f'input size {(input_h, input_w)} is `x+1` and '
                        f'out size {(output_h, output_w)} is `nx+1`')
    if scale_factor is not None:
        h, w = size[0], size[1]
        scale_size = (h*scale_factor, w*scale_factor)
    else:
        scale_size = size

    resize_mode = None
    if mode == 'nearest':
        resize_mode = ops.ResizeNearestNeighbor(size=scale_size, align_corners=align_corners)
    elif mode == 'bilinear':
        resize_mode = ops.ResizeBilinear(size=scale_size, align_corners=align_corners)


    out = resize_mode(input)

    # return resize_mode(input)
    return out
















